import json
from typing import Dict
from openai import AsyncOpenAI
from json_repair import repair_json
from service.chat.chat_manager import ChatManager, ChatState
from service.chat.struct_info import InputsInfo, CompleteInfo, DetailMathInfo, SolveInputs, SolverResults, MathKnowledgeResult
from utils.logger import get_logger
from flask import jsonify


logger = get_logger(__name__)

class ChatServer:
    def __init__(self, config: Dict[str, object]):
        logger.info("Initializing ChatServer")
        self.set_up_model(config)
        self.load_system_prompts(config)
        self.chat_manager = ChatManager.get_instance()

    def set_up_model(self, config: Dict[str, object]):
        self.extractor_model = config["extractor_model"]
        self.main_model = config["main_model"]
        self.math_model = config["math_model"]
        self.client = AsyncOpenAI(api_key=config["api_key"], base_url=config["base_url"])

    def load_system_prompts(self, config: Dict[str, object]):
        self.extract_math_info_prompts = None
        with open(config["prompts_files"]["extract_math_info"], "r") as fr:
            self.extract_math_info_prompts = fr.read()
        
        self.explain_math_prompts = None
        with open(config["prompts_files"]["explain_math_knowledge"], "r") as fr:
            self.explain_math_prompts = fr.read()

        self.math_detail_info_prompts = None
        with open(config["prompts_files"]["detail_math_info"], "r") as fr:
            self.math_detail_info_prompts = fr.read()

        self.math_solve_prompts = None
        with open(config["prompts_files"]["solve_math"], "r") as fr:
            self.math_solve_prompts = fr.read()

    def encode_message(self, system_prompt: str=None, img_content: Dict[str, object] = None, str_content: str= None)->list:
        messages = []
        if system_prompt is not None:
            messages.append({"role": "system", "content": system_prompt})
        
        user_content = []
        if img_content is not None:
            img_data = img_content["img_data"]
            img_type = img_content["img_type"]
            if img_type.endswith("png"):
                user_content.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_data}"}})
            elif img_type.endswith("jpg"):
                user_content.append({"type": "image_url", "image_url": {"url": f"data:image/jpg;base64,{img_data}"}})
        
        if str_content is not None and str_content != "":
            user_content.append({"type": "text", "text": str_content})
        
        if len(user_content) > 0:
            messages.append({"role": "user", "content": user_content})
        
        return messages

    async def extract_math_content(self, session_id: str, input_data: Dict[str, object]) ->Dict[str, object]:
        messages = self.encode_message(self.extract_math_info_prompts, input_data, input_data["text_data"])
        completion = await self.client.chat.completions.create(
            model=self.extractor_model,
            messages=messages,
            temperature=0.,
            top_p=0.,
            n=1
        )
        raw_message = completion.choices[0].message.content
        raw_message = raw_message.replace("```json", "").replace("```", "")
        raw_message = repair_json(raw_message)
        parsed_json = json.loads(raw_message)
        result_value = InputsInfo(parsed_json)
        self.chat_manager.update_current_state(session_id, ChatState.EXTRACT_INFO, result_value)
        return jsonify({
            "id": session_id,
            "state": "EXTRACT_INFO",
            "data": result_value.to_dict()
        }), 200

    async def classfiy_math_info(self, session_id: str) ->Dict[str, object]:
        input_datas: InputsInfo = self.chat_manager.get_current_inputs(session_id)
        messages = self.encode_message(self.math_detail_info_prompts, str_content=input_datas.latex)
        completion = await self.client.chat.completions.create(
            model=self.main_model,
            messages=messages,
            temperature=0.,
            top_p=0.,
            n=1
        )

        raw_message = completion.choices[0].message.content
        raw_message = raw_message.replace("```json", "").replace("```", "")
        raw_message = repair_json(raw_message)
        parsed_json = json.loads(raw_message)
        return_value = DetailMathInfo(parsed_json)
        self.chat_manager.update_current_state(session_id, ChatState.DETAIL_INFO, return_value)
        return jsonify({
            "id": session_id,
            "state": "DETAIL_INFO",
            "data": return_value.to_dict()
        }), 200

    async def solve_math_problem(self, session_id: str) ->Dict[str, object]:
        input_datas: SolveInputs = self.chat_manager.get_current_inputs(session_id)
        messages = self.encode_message(self.math_solve_prompts, str_content=input_datas.to_string())
        completion = await self.client.chat.completions.create(
            model=self.math_model,
            messages=messages,
            temperature=0.,
            top_p=0.,
            n=1
        )

        raw_message = completion.choices[0].message.content
        raw_message = raw_message.replace("```json", "").replace("```", "")
        raw_message = repair_json(raw_message)
        parsed_json = json.loads(raw_message)
        return_value = SolverResults(parsed_json)

        self.chat_manager.update_current_state(session_id, ChatState.SOLVE_PROBLEM, return_value)
        return jsonify({
            "id": session_id,
            "state": "SOLVE_PROBLEM",
            "data": return_value.to_dict()
        }), 200

    async def process_math_knowledge(self, session_id: str) ->Dict[str, object]:
        input_datas: InputsInfo = self.chat_manager.get_current_inputs(session_id)
        messages = self.encode_message(self.explain_math_prompts, str_content=input_datas.latex)
        completion = await self.client.chat.completions.create(
            model=self.main_model,
            messages=messages,
            temperature=0.,
            top_p=0.,
            n=1
        )
        raw_message = completion.choices[0].message.content
        raw_message = raw_message.replace("```json", "").replace("```", "")
        raw_message = repair_json(raw_message)
        parsed_json = json.loads(raw_message)
        return_value = MathKnowledgeResult(parsed_json)
        self.chat_manager.update_current_state(session_id, ChatState.EXPLAIN_KNOWLEDGE, return_value)
        return jsonify({
            "id": session_id,
            "state": "EXPLAIN_KNOWLEDGE",
            "data": return_value.to_dict()
        }), 200

    async def complete(self, session_id: str, reasion: str) -> Dict[str, object]:
        complet_info = CompleteInfo(reasion)
        self.chat_manager.update_current_state(session_id, ChatState.COMPLETE, complet_info)
        
        return jsonify({
            "id": session_id,
            "state": "COMPLETE",
            "data": complet_info.to_dict()
        }), 200

    async def solve(self, session_id: str, input_data:Dict[str, object]) -> Dict[str, object]:
        # 1. 根据session id获取manager信息
        current_state = self.chat_manager.get_current_state(session_id)
        if current_state == ChatState.INIT:
            return await self.extract_math_content(session_id, input_data)
        elif current_state == ChatState.EXTRACT_INFO:
            input_info: InputsInfo = self.chat_manager.get_current_inputs(session_id)
            if input_info.status == "MathKnowledge":
                return await self.process_math_knowledge(session_id)
            elif input_info.status == "MathProblem":
                return await self.classfiy_math_info(session_id)
            elif input_info.status == "NoMath":
                return await self.complete(session_id, "unvalid inputs")
        elif current_state == ChatState.DETAIL_INFO:
            return await self.solve_math_problem(session_id)
        elif current_state == ChatState.SOLVE_PROBLEM:
            return await self.complete(session_id, "problem solved")
        elif current_state == ChatState.EXPLAIN_KNOWLEDGE:
            return await self.complete(session_id, "knowledge explained")
        elif current_state == ChatState.COMPLETE:
            self.chat_manager.reset(session_id)
            return jsonify({
                "id": session_id,
                "state": "RESET",
                "data": ""
            }), 200

    def encode_image(self, image_path):
        import base64
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode("utf-8")

    async def test(self):
        session_id = "test"
        input_data = {
            "text_data": "",
            "img_data": self.encode_image("src_img/002.png"),
            "img_type": "png"
        }

        await self.extract_math_content(session_id, input_data)

